/*
 *Case2:Computing u_t + a u_x = bb u_{xx} = 0 with implicit scheme
 *Author: Zhicheng Yang
 */

static char help[] = "Solves a linear system with KSP.\n\n";

# include <stdio.h>
# include <math.h>
# include <petscksp.h>
# include <stdlib.h>

# define N      100
# define nu     5.0
# define time   1.0
# define a      1
# define bb     0.0475 				/* 这里精确解为u(x,t) = e^{20/19 * x - t} */

Mat A;
Vec x,b;


void initial(double *u,double dx)
{
	int j;
	FILE *fp;

	for (j=0;j<=N;j++)
		*(u+j)=exp(20.0*j*dx/19);
	fp = fopen("initcase2.txt","w");
	for (j = 0;j < N+1;j++)
		fprintf(fp,"%e\n",*(u+j));
	fclose(fp);
}

void output(double *u,double dx)
{
	int j;
	FILE *fp;
	fp=fopen("outputcase2.txt","w");
	for (j=0;j<=N;j++){
		fprintf(fp,"%e\n",*(u+j));
	}
	fclose(fp);
	fp=fopen("case1Config.txt","a");
	fprintf(fp,"Using Implicit Method Computing equ u_t + u_x = 0.0475 u_{xx} = 0\n with nu = %e, time = %e\n",nu,time);
	fclose(fp);
}

/* L^2 and L^{\infty} error */

void outputerror(double *u,double dx,int step)
{
	int    j;
	double err,err2,e;
	double exact[N+1];
	FILE   *fp;

	err = 0.0; err2 = 0.0; e = 0.0;
	for (j=0;j<N+1;j++)
		*(exact+j)=exp(20.0*j*dx/19-time);
	for (j=0;j<N+1;j++){
		e=fabs(*(exact+j)-*(u+j));
		if (e>err)
			err=e;
		err2+=e*e;
		//printf("%e\n",e);
	}
	err2=sqrt(err2*dx);

	fp=fopen("exactcase2.txt","w");
	for (j = 0;j<N+1;j++)
		fprintf(fp,"%e\n",*(exact + j));
	fclose(fp);
	
	fp = fopen("case2err.txt","a");
	fprintf(fp,"Case2: u_t + u_x = 0.0475 u_{xx}\n with time = %f, dx=%e, nu= %e.\n",time,dx,nu);
	fprintf(fp,"Compute Steps:%d\n",step);
	fprintf(fp,"L^{\\infty} error:%e\n",err);
	fprintf(fp,"L^2 error:%e\n\n\n",err2);
	fclose(fp);
}


void constrmat(double coura,double newnu)
{
	int                 *row,*colum;
	PetscInt            j,nz;                 
	PetscInt            *nnz;
	double              a_n1,a_0,a_p1;
 
	a_n1 = -coura/2 - newnu;
	a_0 = 1 + 2*newnu;
	a_p1 = coura/2 - newnu;

	nz=3;nnz=PETSC_NULL;
	row = (int *) malloc(sizeof(int));
	colum = (int *) malloc(sizeof(int));
	MatCreateSeqAIJ(PETSC_COMM_SELF,N-1,N-1,nz,nnz,&A);
	for (j = 0;j < N-1;j++){
		if (j!=0 && j!=N-2){
			*row = j;
			*colum = *row - 1;
			MatSetValues(A,1,row,1,colum,&a_n1,INSERT_VALUES);
			*colum = *row;
			MatSetValues(A,1,row,1,colum,&a_0,INSERT_VALUES);
			*colum = *row + 1;
			MatSetValues(A,1,row,1,colum,&a_p1,INSERT_VALUES);
		}
		else{
			if (j == 0){
				*row = j;
				*colum = *row;
				MatSetValues(A,1,row,1,colum,&a_0,INSERT_VALUES);
				*colum = *row + 1;
				MatSetValues(A,1,row,1,colum,&a_p1,INSERT_VALUES);
			}
			else{
				*row = j;
				*colum = *row - 1;
				MatSetValues(A,1,row,1,colum,&a_n1,INSERT_VALUES);
				*colum = *row;
				MatSetValues(A,1,row,1,colum,&a_0,INSERT_VALUES);
			}
		}
	}
	free(row);
	free(colum);
	MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
	MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
}
int main(int argc, char ** args)
{
	PetscInt            j,nz,step;
	PetscInt            *nnz;
	PetscReal           t;
	int                 *row,*colum;
	PetscErrorCode      ierr;
	KSP                 ksp;
	PC                  pc;
	PetscMPIInt         size,rank;
	double              a_n1,a_0,a_p1,coura,coura0,newnu,dx,dt; /* In scheme, a_n1 u_{j-1}^{n+1} + a_0 u_{j}^{n+1} + a_p1 u_{j+1}^{n+1} = u_{j}^{n} */
	double              u[N+1],u0[N+1];
	int                 index[N-1];
	
	PetscInitialize(&argc,&args,(char *)0,help);
	ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRQ(ierr);
	ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRQ(ierr);

	dx = 1.0/N;dt = nu*dx*dx/bb;t = 0.0;
	coura = a*dt/dx; 
	a_n1 = -coura/2 - nu;
	a_0 = 1 + 2*nu;
	a_p1 = coura/2 - nu;
	step = 0;
	
	
	nz=3;nnz=PETSC_NULL;
	row = (int *) malloc(sizeof(int));
	colum = (int *) malloc(sizeof(int));
	MatCreateSeqAIJ(PETSC_COMM_SELF,N-1,N-1,nz,nnz,&A);
	for (j = 0;j < N-1;j++){
		if (j!=0 && j!=N-2){
			*row = j;
			*colum = *row - 1;
			MatSetValues(A,1,row,1,colum,&a_n1,INSERT_VALUES);
			*colum = *row;
			MatSetValues(A,1,row,1,colum,&a_0,INSERT_VALUES);
			*colum = *row + 1;
			MatSetValues(A,1,row,1,colum,&a_p1,INSERT_VALUES);
		}
		else{
			if (j == 0){
				*row = j;
				*colum = *row;
				MatSetValues(A,1,row,1,colum,&a_0,INSERT_VALUES);
				*colum = *row + 1;
				MatSetValues(A,1,row,1,colum,&a_p1,INSERT_VALUES);
			}
			else{
				*row = j;
				*colum = *row - 1;
				MatSetValues(A,1,row,1,colum,&a_n1,INSERT_VALUES);
				*colum = *row;
				MatSetValues(A,1,row,1,colum,&a_0,INSERT_VALUES);
			}
		}
	}
	free(row);
	free(colum);
	MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
	MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
	initial(u0,dx);
	
	coura0=coura;
	newnu=nu;
	while (t < time){
		if (t+dt>time){
			dt = time - t;
			MatDestroy(A);
			newnu = bb*dt/(dx*dx);
			coura0 = a*dt/dx;
			constrmat(coura0,newnu);
		}
		t+=dt;
		step++;
		VecCreateSeq(PETSC_COMM_SELF,N-1,&b);
		VecDuplicate(b,&x);
		for (j = 0; j<N-1; j++){
			VecSetValue(b,j,*(u0+j+1),INSERT_VALUES);
		}
		VecSetValue(b,0,*(u0+0)*(coura0/2+newnu),ADD_VALUES);
		VecSetValue(b,N-2,*(u0+N)*(-coura0/2+newnu),ADD_VALUES);
		VecAssemblyBegin(b);	
		VecAssemblyEnd(b);
		
		/* Create ksp object, and set CG, BJacobi as the compute method and preconditioner respectively */
		KSPCreate(PETSC_COMM_WORLD,&ksp);
		KSPSetOperators(ksp,A,A,DIFFERENT_NONZERO_PATTERN);
		KSPGetPC(ksp,&pc);
		PCSetType(pc,PCBJACOBI);
		KSPSetType(ksp,KSPBICG);
		KSPSetTolerances(ksp,1.e-6,PETSC_DEFAULT,PETSC_DEFAULT,PETSC_DEFAULT);
		KSPSetFromOptions(ksp);
		//solve
		KSPSolve(ksp,b,x);
		//t=MPI_Wtime()-t;              /* return compute time */
		//ierr = KSPView(ksp,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
		
		for (j=0;j<N-1;j++)
			*(index+j) = j;
		
		VecGetValues(x,N-1,index,(u+1));
		*(u) = exp(-t);
		*(u + N) = exp(20.0/19 - t);
		for (j = 0;j<N+1;j++)
			*(u0 + j) = *(u + j);
		VecDestroy(x);VecDestroy(b);
	}
	
	output(u,dx);
	outputerror(u,dx,step);
	//free
	MatDestroy(A);
	ierr = KSPDestroy(ksp);CHKERRQ(ierr);
	
	PetscFinalize();
	printf("t = %e\nstep = %d\n",t,step);
	return 0;
}
